'''
需要输入的参数包括
input（N, 7168）: 输入隐藏状态张量 通常来自上一层Transformer块的输出
作用：作为MLA预处理的起始输入，经过RMS Norm1量化后用于dqkv投影。

N：序列长度 即当前处理token数量
作用：用于重塑张量维度、RoPE应用和缓存槽位映射。

1. RMS Norm1 相关参数（用于input1 -> dqkv输入的量化）
gamma1（7168, ）：RMS Norm1的gamma参数，用于归一化后的缩放。
作用：在RMS Norm计算中乘以归一化后的输入，实现层归一化。

beta1（7168, ）：RMS Norm1的beta参数（bias），用于归一化后的偏移。
作用：在RMS Norm量化计算中添加偏移，实现带偏的归一化。

quant_scale1（1,  ）：RMS Norm1的量化缩放因子（scale），用于将浮点输出映射到int8范围。
作用：用于量化RMS Norm输出：output = (norm + beta) * scale + offset。

quant_offset1（1, ）：RMS Norm1的量化偏移（offset），用于int8量化的零点调整。
作用：确保量化后的值在[QUANTMIN, QUANTMAX]（-128~127）范围内。

2. dpkv 投影相关参数
wdokv（2112, 7168）：dqkv权重矩阵（weight for down-project qkv），int8量化。
作用：用于MatMul计算qkv投影：mm_out = input @ wdqkv.T + bias。

de_scale1（2112, ）：dqkv的反量化缩放因子（dequantize scale），float32。
作用：用于反量化MatMul输出：mm_out * de_scale，实现从int32到float16的转换。

bias1（1， 2112）：dqkv偏置（bias），int32类型。
作用：添加到MatMul输出中，提供线性变换的偏移。
说明：dpkv输出被拆分为key_norm（512维）、key_rope（64维）、value_norm（1536维）

3. RMS Norm2 相关参数（用于value_norm -> uq输入的量化）
gamm2（1536, ）：RMS Norm2的gamma参数（scale factor）。
作用：缩放value_norm的RMS Norm输出。

beta2（1536, ）：RMS Norm2的beta参数（bias）。
作用：偏移value_norm的RMS Norm输出。

quant_scale2（1, ）：RMS Norm2的量化缩放因子。
作用：量化RMS Norm2输出到int8。

quant_offset2（1, ）：RMS Norm2的量化偏移
作用：调整量化零点。

4. uq 投影相关参数
wuq（head_num * 192, 1536）：uq权重矩阵（weight for up-project q），int8量化。
作用：用于MatMul计算uq投影：mm_out = value_norm @ wuq.T + bias。

de_scale2（head_num * 192）：uq的反量化缩放因子，float32。
作用：反量化uq MatMul输出到float16。

bias2（1， head_num * 192）：uq偏置，int32。
作用：添加到uq MatMul输出中。

5. RoPE 和 RMS Norm3 相关参数
gamma3（512, ）：RMS Norm3的gamma参数，用于key_norm的归一化缩放。
作用：应用于key_norm的RMS Norm计算。

cos1（N, 64）：RoPE的余弦位置编码（cosine），针对每个token和rope_hidden_size=64维。
作用：用于Q/K的RoPE旋转：q/k * cos + rotate(q/k) * sin。

sin1（N, 64）：RoPE的正弦位置编码（sine）。
作用：同上，扩展到head_num维度后应用。

5. K投影和缓存相关参数
wuk（head_num, 128, 512）：K投影权重（weight for key projection），用于q_base到q的线性变换。
作用：通过BMM（batch matmul）计算：q_base @ wuk，实现head-wise投影。

key_cache（block_num, block_size, 1, 576）：K缓存张量（key cache），block_num=192（总块数）。
作用：存储过去的key值，支持KV缓存机制，用于加速自回归生成。

slot_mapping（N, ）：槽位映射（slot mapping），int32，表示每个token在缓存中的位置索引。
作用：用于Reshape & Cache：将当前key插入到key_cache的指定槽位。

head_num：注意力头数（number of heads），如16/32/64/128。
作用：用于拆分多头注意力、RoPE扩展和uq/wuk维度。

block_size：缓存块大小（block size），如128/256。
作用：定义每个缓存块的token容量，用于slot_mapping计算block_index和offset。

data_type：数据类型（dtype），支持float16或bfloat16。
作用：指定输出和中间计算的精度，影响dequantize处理。
'''
import torch
import torch_npu
import pytest
import numpy as np
import random
from mlapo_torch import MLAPO
from mlapo_triton import mlapo_triton

N = [1024]
HEAD_NUM = [16, 32, 64, 128]
BLOCK_SIZE = [128, 256]
DTYPE = [torch.float16]
BLOCK_NUM = [192]
HEAD_DIM = [576]
HIDDEN_SIZE = [7168] # 模拟数据
HIDDEN_SIZE_WDQKV = [2112]
HIDDEN_SIZE_WDQ = [1536]
HIDDEN_SIZE_ROPE_K = [64]
HIDDEN_SIZE_WUQ_HEAD = [192]
HIDDEN_SIZE_WUK_HEAD = [512]
CACHE_MODE = [0]
epsilon = 1e-6
DEVICE = "npu"
seed = 12

@pytest.mark.parametrize("tokens", N)
@pytest.mark.parametrize("head_num", HEAD_NUM)
@pytest.mark.parametrize("block_size", BLOCK_SIZE)
@pytest.mark.parametrize("block_num", BLOCK_NUM)
@pytest.mark.parametrize("head_dim", HEAD_DIM)
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZE)
@pytest.mark.parametrize("hidden_size_wdqkv", HIDDEN_SIZE_WDQKV)
@pytest.mark.parametrize("hidden_size_wdq", HIDDEN_SIZE_WDQ)
@pytest.mark.parametrize("hidden_size_rope_k", HIDDEN_SIZE_ROPE_K)
@pytest.mark.parametrize("hidden_size_wuq_head", HIDDEN_SIZE_WUQ_HEAD)
@pytest.mark.parametrize("hidden_size_wuk_head", HIDDEN_SIZE_WUK_HEAD)
@pytest.mark.parametrize("cache_mode", CACHE_MODE)
@pytest.mark.parametrize("data_type", DTYPE)
@torch.inference_mode()
def test_mlapo(
    tokens: int,
    head_num: int,
    block_size: int,
    block_num: int,
    head_dim: int,
    hidden_size: int,
    hidden_size_wdqkv: int,
    hidden_size_wdq: int, 
    hidden_size_rope_k: int,
    hidden_size_wuq_head: int,
    hidden_size_wuk_head: int,
    cache_mode: int,
    data_type: torch.dtype
):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    # 设置使用的机器号
    torch.npu.set_device(7)
    # 元数据
    # tokens = 1024
    # head_num_q = random.choice([32, 128, 16, 64])
    # block_size = random.choice([128, 256])
    # block_num = 192
    # data_type = random.choice([torch.float16, torch.bfloat16])
    
    # hidden_size = 7168
    ## d表示降维，u表示升维
    # hidden_size_wdqkv = 2112
    # hidden_size_wdq = 1536
    hidden_size_wdkv = hidden_size_wdqkv - hidden_size_wdq  ## 576
    # hidden_size_rope_k = 64
    hidden_size_rms_kv = hidden_size_wdkv - hidden_size_rope_k  ## 512
    # hidden_size_wuq_head = 192
    hidden_size_rope_q_head = hidden_size_rope_k
    hidden_size_nope_q_head = hidden_size_wuq_head - hidden_size_rope_q_head
    # hidden_size_wuk_head = 512

    hidden_size_output_kv = hidden_size_rope_k + hidden_size_rms_kv  ## 576
    hidden_size_output_q = hidden_size_rope_q_head + hidden_size_wuk_head  ## 576

    # epsilon = 1e-6
    # cache_mode = 0

    # 生成模拟输入（实际使用时替换为真实数据）
    input = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(tokens, hidden_size))).to(data_type)
    gamma1 = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(hidden_size))).to(data_type)
    beta1 = torch.from_numpy(np.random.randint(-2, 2, (hidden_size)).astype(np.float16)).to(data_type)
    quant_scale1 = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(1))).to(data_type)
    quant_offset1 = torch.from_numpy(np.random.uniform(-128.0, 127.0, size=(1))).to(torch.int8)
    # T
    # wdqkv = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(hidden_size, hidden_size_wdqkv))).to(torch.int8)
    wdqkv = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(hidden_size_wdqkv, hidden_size))).to(torch.int8)
    de_scale1 = torch.rand((hidden_size_wdqkv), dtype=torch.float32) / 1000
    bias1 = torch.from_numpy(np.random.randint(-10, 10, (1, hidden_size_wdqkv)).astype(np.int32)).to(torch.int32)

    gamma2 = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(hidden_size_wdq))).to(data_type)
    beta2 = torch.from_numpy(np.random.randint(-2, 2, (hidden_size_wdq)).astype(np.float16)).to(data_type)
    quant_scale2 = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(1))).to(data_type)
    quant_offset2 = torch.from_numpy(np.random.uniform(-128.0, 127.0, size=(1))).to(torch.int8)
    # T
    # wuq = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(hidden_size_wdq, head_num * hidden_size_wuq_head))).to(torch.int8)
    wuq = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(head_num * hidden_size_wuq_head, hidden_size_wdq))).to(torch.int8)
    de_scale2 = torch.rand((head_num * hidden_size_wuq_head), dtype=torch.float32) / 1000
    bias2 = torch.from_numpy(np.random.randint(-10, 10, (1, head_num * hidden_size_wuq_head)).astype(np.int32)).to(torch.int32)

    gamma3 = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(hidden_size_rms_kv))).to(data_type)
    cos1 = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(tokens, hidden_size_rope_k))).to(data_type)
    sin1 = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(tokens, hidden_size_rope_k))).to(data_type)
    cos2 = cos1
    sin2 = sin1
    wuk = torch.from_numpy(np.random.uniform(-2.0, 2.0, size=(head_num, hidden_size_nope_q_head, hidden_size_wuk_head))).to(data_type)
    key_cache = torch.from_numpy(np.random.uniform(-1.0, 1.0, size=(block_num, block_size, 1, hidden_size_output_kv))).to(data_type)
    slot_mapping = torch.from_numpy(np.random.choice(192 * 128, tokens, replace=False).astype(np.int32)).to(torch.int32)

    # 调用 Golden 计算
    mlapo = MLAPO()
    golden_outputs = mlapo.mla_preprocess_calc_(
        input, 
        gamma1, 
        beta1, 
        quant_scale1, 
        quant_offset1, 
        wdqkv, 
        de_scale1, 
        bias1,
        # 
        gamma2, 
        beta2, 
        quant_scale2, 
        quant_offset2, 
        wuq, 
        de_scale2, 
        bias2, 
        # 
        gamma3,
        cos1, 
        sin1, 
        cos2, 
        sin2, 
        wuk, 
        key_cache, 
        slot_mapping,
        # 
        tokens=tokens, 
        head_num_q=head_num, 
        block_size=block_size, 
        data_type=data_type, 
        cache_mode = cache_mode, 
        # 
        epsilon = epsilon, 
        hidden_size = hidden_size, 
        hidden_size_wdq = hidden_size_wdq, 
        hidden_size_rope_q_head = hidden_size_rope_q_head, 
        hidden_size_rope_k = hidden_size_rope_k
    )
    actual_output = mlapo_triton(
        input, gamma1, beta1, quant_scale1, quant_offset1, wdqkv, de_scale1, bias1,
        gamma2, beta2, quant_scale2, quant_offset2, wuq, de_scale2, bias2, gamma3,
        cos1, sin1, cos2, sin2, wuk, key_cache, slot_mapping,
        Tokens=tokens, head_num=head_num, block_size=block_size, data_type=data_type, cache_mode=cache_mode
    )
    atol, rtol = 1.5e-1, 1.5e-1
    torch.testing.assert_close(golden_outputs[0], actual_output[0], rtol=rtol, atol=atol)
    torch.testing.assert_close(golden_outputs[1], actual_output[1], rtol=rtol, atol=atol)
    # match_rate = torch.allclose(golden_outputs[0].float(), actual_output[0].float(), atol=1e-2) # int8转回float进行比较
    # print(f"与 PyTorch 匹配: {match_rate}")
    # if not match_rate:
    #     mismatch_count = (golden_outputs[0] != actual_output[0]).sum().item()
    #     print(f"不匹配数量: {mismatch_count} / {golden_outputs[0].numel()} ({mismatch_count / golden_outputs[0].numel() * 100:.4f}%)")

if __name__ == "__main__":
    test_mlapo(1024, 16, 128, 192, 576, 7168, 2112, 1536, 64, 192, 512, 0, torch.float16)